import gradio as gr
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Qdrant
from langchain.memory import ConversationSummaryMemory
from langchain_zhipu import ChatZhipuAI, ZhipuAIEmbeddings
from langchain.chains import ConversationalRetrievalChain
from WriteDocuments import FAISSService
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain
from langchain.chains.retrieval import create_retrieval_chain


class ChatbotWithRetrieval:
    def __init__(self, base_dir):
        model = ChatZhipuAI(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"
                            , model="glm-4", verbose=True)

        model.do_sample = False  # 温度设置为0，结果随机性 ghbnm
        faiss = FAISSService(base_dir)
        vectorstore = faiss.vectorstore

        prompt = ChatPromptTemplate.from_template("""Answer the following questions based on context and historical questions and answers,
        Please answer in Chinese:

        <context>
        {context}
        </context>
        
        <history>
        {history}
        </history>
        Question: {input}
        
        Priority should be given to answering with the given information. If you do not know the answer, please use your own knowledge base to answer. If you do not know, please answer truthfully and do not guess.
Please try to be concise and concise in your answers""")
        document_chain = create_stuff_documents_chain(model, prompt)

        retriever = vectorstore.as_retriever()
        retrieval_chain = create_retrieval_chain(retriever, document_chain)

        self.store = faiss
        # 向量数据库
        self.vectorstore = vectorstore

        # 初始化LLM
        self.llm = model

        # 初始化Memory
        self.memory = ""
        # 初始化对话历史
        self.conversation_history = ""
        self.qa = retrieval_chain
